#include "Cylinder.h"
#include "Mat6x6.h"
#include "PrescribedMotion.h"
#include "ABAData.h"
#include "EOMData.h"
#include "RigidBodyNode.h"

using namespace std;
using namespace RSIM;

///////////////////////////////////////////////////////

int Cylinder::getJointDOF() const{return 2;}

///////////////////////////////////////////////////////

void Cylinder::updRelVelocity(const State& state, Vect3& w, Vect3& v) const{
        
        // Angular velocity of origin
        w(2) = state.getUForNode(this->NodeIndex_)[0];
        
        // linear velocity of origin
        v(2) = state.getUForNode(this->NodeIndex_)[1];
}

///////////////////////////////////////////////////////

void Cylinder::calc_H_x_UDot(const State& state, double *h_x_udot) const{
        h_x_udot[2] = state.getUDotForNode(this->NodeIndex_)[0]; 
        h_x_udot[5] = state.getUDotForNode(this->NodeIndex_)[1];
}

///////////////////////////////////////////////////////

void Cylinder::setDefaultQ(double *q)const{
        q[0] = 0.0;	q[1] = 0.0;	
}

///////////////////////////////////////////////////////

void Cylinder::setDefaultU(double *u)const{
        u[0] = 0.0;	u[1] = 0.0;
}

///////////////////////////////////////////////////////

void Cylinder::calcQDot(State& state)const{
        state.wgetQDotForNode(this->NodeIndex_)[0] = state.wgetUForNode(this->NodeIndex_)[0];
        state.wgetQDotForNode(this->NodeIndex_)[1] = state.wgetUForNode(this->NodeIndex_)[1];
}

//////////////////////////////////////////////////////

Cylinder::Cylinder(const int& NodeIndex):Joint(CylinderJoint){
        this->NodeIndex_ = NodeIndex;
}

//////////////////////////////////////////////////////

Cylinder* Cylinder::New(const int& NodeIndex){
        return new Cylinder(NodeIndex);
}

///////////////////////////////////////////////////////

void Cylinder::printJointType() const{
        cout<<"Joint Type=Cylinder\n";
}

///////////////////////////////////////////////////////

void Cylinder::updJointTransform(const double *q,
                                const RigidBodyNode *Parent,
                                const int& pChildID,
                                double *betak){
        
        // Update joint transformation between frames M and F
        this->T_MF_.wgetR().setRotationZAxis(q[0]);
        
        // Update the length of vector v_MF
        this->T_MF_.wgetP().wPtr()[2]=q[1];
        
        // update betak
        this->updBetaK(q,Parent,pChildID,betak);	
}

///////////////////////////////////////////////////////

void Cylinder::updBetaK(	const double *q,
                                const RigidBodyNode *Parent,
                                const int& pChildID,
                                double *betak)const{
        if(Parent){
                // Calculate betak in F basis of parent
                const double *T_FM_P_Ptr = Parent->data_->getBodyTransformFM()[pChildID].P().Ptr();
                const double *T_FM_R_Ptr = Parent->data_->getBodyTransformFM()[pChildID].R().Ptr();
                
                betak[0]=T_FM_P_Ptr[0];
                betak[1]=T_FM_P_Ptr[1];
                betak[2]=T_FM_P_Ptr[2];
                
                betak[0] += T_FM_R_Ptr[2]*q[1];
                betak[1] += T_FM_R_Ptr[5]*q[1];
                betak[2] += T_FM_R_Ptr[8]*q[1];
        }
}

///////////////////////////////////////////////////////

void Cylinder::updVelocity(	const State& state, 
                                        const RigidBodyNode *Parent,
                                        const Rotation& prF_C_F, 
                                        const double *betak,
                                        Vect3& w,
                                        Vect3& v) const {
                                        
        const Vect3& pr_w = Parent->data_->getAngularVelocity();
        const Vect3& pr_v = Parent->data_->getLinearVelocity();
        
        // 1. F_C_prF*pr_w is the absolute angular velocity
        // of the parent expressed in the joint's F frame
        // 2. Vect3(0,0,U_[0]) is the angular velocity of this body
        // which is already expressed in the joint's F frame
        double *wptr = w.wPtr();
        FastOps::transpose_mult_mat33_vec3(prF_C_F.Ptr(),pr_w.Ptr(),wptr);
        wptr[2] += state.getUForNode(this->NodeIndex_)[0];
        
        double v_tmp[3];
        
        // omega x r
        FastOps::cross(pr_w.Ptr(),betak,v_tmp);
        
        // v + omega x r
        v_tmp[0] += pr_v(0);
        v_tmp[1] += pr_v(1);
        v_tmp[2] += pr_v(2);
        
        // convert v + omega x r to F basis of this body
        FastOps::transpose_mult_mat33_vec3(prF_C_F.Ptr(), v_tmp, v.wPtr());
        
        // Velocity = v + omega x r + qdot
        v.wPtr()[2] += state.getUForNode(this->NodeIndex_)[1];
}

///////////////////////////////////////////////////////


void Cylinder::calcAKT(		const RigidBodyNode *Parent, 
                                const RigidBodyNode *Body,
                                const int& ChildID,
                                const double *betak,
                                State& state,
                                Vect6& sAkt) const{
        this->calcQDot(state);
        double *UDot = state.wgetUDotForNode(this->NodeIndex_);
        if(Parent){
                const double *U = state.getUForNode(this->NodeIndex_);
                
                // AKT is calculated in F basis of this body
                double *akt_ptr = sAkt.wPtr();
                
                const double *prFCF_Ptr = Body->data_->getRotationFCF().Ptr();
                
                double s_betak[3], ang_vel[3];
                
                // Angular velocity of parent expressed in F of parent
                const double *parent_ang_ptr = Parent->data_->getAngularVelocity().Ptr();
                
                // express angular velocity and betak in F basis of this body
                FastOps::transpose_mult_mat33_vec3(prFCF_Ptr,parent_ang_ptr,ang_vel);
                FastOps::transpose_mult_mat33_vec3(prFCF_Ptr,betak,s_betak);
                
                // cross(w^k, w^{k+1}_{k+1}u^{k+1})
                akt_ptr[0] = U[0]*ang_vel[1];
                akt_ptr[1] = -U[0]*ang_vel[0];
                akt_ptr[2] = 0.0;
                
                // cross( w^{k-1}, cross(w^{k-1}, \beta^{k-1}) ) {expressed in F of this body}
                FastOps::wxwxr(ang_vel, s_betak, akt_ptr+3);				
                
                // 2*cross(w^k, v^{k+1}_{k+1}u^k+1)
                akt_ptr[3] += 2*U[1]*ang_vel[1];
                akt_ptr[4] += -2*U[1]*ang_vel[0];
                        
                // In this case, Akt = Akt + p^k_k \dot{u}^k
                if(this->isMotionPrescribedAndActive_){
                        // Update QDot_ cache from prescribed motion
                        this->PM_->setUDot(state.getSysData(),UDot);
                        akt_ptr[2] = UDot[0];
                        
                        // The += sign can be replaced by =. akt_ptr[5] should be zero until this point.
                        // need to check this at some point.
                        akt_ptr[5] += UDot[1];
                }
        } // if(Parent)
        else{
                if(this->isMotionPrescribedAndActive_){
                        this->PM_->setUDot(state.getSysData(),UDot);
                }
        }
}

///////////////////////////////////////////////////////

void Cylinder::sweepBackwards(	const bool& isLeaf, 
                                const int& ChildID,
                                const Vect6& sAkt,
                                const Rotation& FCF,					    
                                const EOMData *eom,
                                const double *betak,
                                RigidBodyNode *Parent,
                                ABAData *aba) const{
        
        Mat6x6 *Ik3 = aba->wgetIk3();
        Vect6 *Fk3 = aba->wgetFk3();		
        
        if(!isMotionPrescribedAndActive_){		
                this->calc_Triang_x_Ik3_TransZ(Ik3->Ptr(), (Ik3+1)->wPtr());
                this->calc_Triang_x_cFk_TransZ(Ik3->Ptr(), Fk3->Ptr(), (Fk3+1)->wPtr());
                
                if(Parent){
                        // cFk = Fk3 - Ik3*A^{k-1}_t
                        this->calc_cFk((Fk3+1)->Ptr(),(Ik3+1)->Ptr(),sAkt.Ptr(),aba->wgetcFk()->wPtr());
                        
                        double triang_ik3[36],triang_fk3[6];

                        this->calc_Triang_x_Ik3_RotZ((Ik3+1)->Ptr(), triang_ik3);
                        this->calc_Triang_x_cFk_RotZ((Ik3+1)->Ptr(),aba->getcFk()->Ptr(),triang_fk3);

                        double sIk3[36], sFk3[6];

                        FastOps::spatial_basis_shift_mat(FCF.Ptr(), triang_ik3, sIk3);
                        FastOps::spatial_basis_shift_vec(FCF.Ptr(), triang_fk3, sFk3);

                        double Ik3Child[36], Fk3Child[6];

                        this->calc_skbeta_x_Ik3_x_tskbeta(sIk3, betak, Ik3Child);
                        this->calc_skbeta_x_cFk(sFk3, betak, Fk3Child);

                        Parent->data_->wgetABAData()->wgetIk3()[0].add(Ik3Child);
                        Parent->data_->wgetABAData()->wgetFk3()[0].add(Fk3Child);
                        
                }// If (Parent)		
        } // if (!isMotionPrescribedAndActive)
        else{
                if(Parent){
                        // cFk = Fk3 - Ik3*A^{k-1}_t
                        this->calc_cFk(Fk3->Ptr(),Ik3->Ptr(),sAkt.Ptr(),aba->wgetcFk()->wPtr());
                        
                        double sIk3[36], sFk3[6];
                        FastOps::spatial_basis_shift_mat(FCF.Ptr(), Ik3->Ptr(), sIk3);
                        FastOps::spatial_basis_shift_vec(FCF.Ptr(), aba->getcFk()->Ptr(), sFk3);
                        
                        double Ik3Child[36], Fk3Child[6];

                        this->calc_skbeta_x_Ik3_x_tskbeta(sIk3, betak, Ik3Child);
                        this->calc_skbeta_x_cFk(sFk3, betak, Fk3Child);
                        
                        Parent->data_->wgetABAData()->wgetIk3()[0].add(Ik3Child);
                        Parent->data_->wgetABAData()->wgetFk3()[0].add(Fk3Child);
                }
        }
}

/////////////////////////////////////////////////////////////////////////


void Cylinder::sweepForward(	const RigidBodyNode *Parent,
                                const int& ChildID,
                                const Vect6& sAkt,
                                const Rotation& FCF,
                                const ABAData *aba,
                                const double *betak,
                                EOMData *eom,
                                State& state
                                ){	
        
        double *UDot = state.wgetUDotForNode(this->NodeIndex_);
        const Vect6 *Fk3 = aba->getFk3();
        const Mat6x6 *Ik3 = aba->getIk3();
        
        Vect6 *Ak1 = eom->wgetAk1();
        double *ak1_ptr = Ak1->wPtr();
        
        if(Parent){
                const double *Parent_Ak1_Ptr = Parent->data_->getEOMData()->getAk1()->Ptr();
                if(!isMotionPrescribedAndActive_){
                        {
                                double ak1_tmp[6];				
                                FastOps::transpose_skbeta_x_V(Parent_Ak1_Ptr, betak, ak1_tmp);
                                FastOps::transpose_spatial_basis_shift_vec(FCF.Ptr(),ak1_tmp, ak1_ptr);
                        }
                        this->calc_UDot_RotZ((Ik3+1)->Ptr(),ak1_ptr,aba->getcFk()->Ptr(),UDot);			
                        
                        Ak1->add(sAkt.Ptr());
                        ak1_ptr[2] += UDot[0];
                        
                        this->calc_UDot_TransZ(Ik3->Ptr(),ak1_ptr,Fk3->Ptr(),UDot+1);
                        ak1_ptr[5] += UDot[1];
                }//if(!isMotionPrescribedAndActive_)
                else{
                        double Ak1p1[6];
                        // Ak1p1 (0:5)= transpose(skbeta)*A^{parent}_1 
                        // In this case, sAkt already includes p^k_k \dot{u}^k
                        FastOps::transpose_skbeta_x_V(Parent_Ak1_Ptr, betak,Ak1p1);
                        
                        // Ak1 = Ak1p1 (0:5) exprssed in F of this body
                        FastOps::transpose_spatial_basis_shift_vec(FCF.Ptr(),Ak1p1,eom->wgetAk1()->wPtr());
                        eom->wgetAk1()->add(sAkt.Ptr());
                }
        }//if(Parent)
        else{
                if(!isMotionPrescribedAndActive_){			
                        UDot[0] = (Fk3+1)->Ptr()[2]/(Ik3+1)->Ptr()[14];
                        UDot[1]=(Fk3->Ptr()[5]-UDot[0]*Ik3->Ptr()[32])/Ik3->Ptr()[35];
                        
                        ak1_ptr[0]=0.0;ak1_ptr[1]=0.0; ak1_ptr[2]=UDot[0];  
                        ak1_ptr[3]=0.0;ak1_ptr[4]=0.0; ak1_ptr[5]=UDot[1];											
                }
                else{
                        // Calculate Ak1_ in F frame of this body
                        double *ak1_ptr = eom->wgetAk1()->wPtr();
                        ak1_ptr[0]=0.0; ak1_ptr[1]=0.0; ak1_ptr[2]=UDot[0];
                        ak1_ptr[3]=0.0; ak1_ptr[4]=0.0; ak1_ptr[5]=UDot[1];
                }		
        }
        // calculate the spatial constraint force
        this->calc_m_cFk(aba->getFk3()[0].Ptr(), aba->getIk3()[0].Ptr(), ak1_ptr, eom->wgetFkc1()->wPtr());
}